from until.untils import *
from openai import OpenAI
from neo4j import GraphDatabase

class Constant:
    NEO4J_USER = "neo4j"
    NEO4J_PASSWORD = "12345678"
    NEO4J_URI = "bolt://localhost:7687"

    DEEPSEEK_API_KEY = "sk-DEEPSEEK_API_KEY"
    DEEPSEEK_BASE_URL = "https://api.deepseek.com"

class Neo4jClient:
    """
    封装Neo4j驱动类
    """
    def __init__(self,username:str,password:str,uri:str="bolt://localhost:7687"):
        self.driver=GraphDatabase.driver(uri, auth=(username, password))

class OpenAIClient():
    """
    封装OpenAiClient类
    """
    def __init__(self,api_key,base_url):
        self.client=OpenAI(api_key=api_key,base_url=base_url)

class ExtractAndCreate:
    """
    执行提取上传等任务类
    """
    def __init__(self,openai_client:OpenAIClient,neo4j_client:Neo4jClient,filePath,showLog):
        """
        初始化类
        :param openai_client: openai类
        :param neo4j_client: neo4j类
        :param filePath: 被提取的文本文件地址
        :param showLog: 是否展示信息
        """
        self.client=openai_client.client
        self.neo4j_client=neo4j_client
        self.filePath=filePath
        self.showLog=showLog

    def __getFileText(self):
        """
        获取文本文件内容
        :return: 文本文件内容
        """
        return open(self.filePath, "r", encoding="utf-8").read()

    def extractFileTextToStrTuple(self):
        """
        提取文本内容中的关系并提取为元组格式字符串
        :return: 元组格式字符串
        """
        client = self.client
        text = self.__getFileText()
        tbl = CreateStrTupleByLLM(client, self.showLog)
        strTuple = tbl.getTuple(text)
        return strTuple

    def strTupleToTuple(self,strTuple: str, splitCode: str = ";"):
        """
        字符串类型的元组转成元组
        :param strTuple: 字符串类型的元组
        :param splitCode: 分割符
        :return:元组列表
        """
        # LLM可能有返回错误，所以需要处理一下
        try:
            dataTuple =DataTuple.getTupleListFromVar(splitCode, strTuple)
        except:
            dataTuple =DataTuple.getTupleListFromVar(splitCode, ";\n")
        return dataTuple

    def createNodeAndRelationToNeo4j(self, dataTuple):
        """
        创建节点和关系
        :param dataTuple: 需要创建的元组和关系的元组列表
        :return: 无返回值
        """
        ttn = TupleToNeo4j(self.neo4j_client.driver)
        ttn.createRelationsAndNodes(dataTuple)


